from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import pycocotools.coco as coco
from pycocotools.cocoeval import COCOeval
import numpy as np
import json
import os

import torch.utils.data as data

class DOTA(data.Dataset):
  num_classes = 9
  default_resolution = [1024, 1024]
  # mean = np.array([0.40789654, 0.44719302, 0.47026115],
  #                  dtype=np.float32).reshape(1, 1, 3)
  # std  = np.array([0.28863828, 0.27408164, 0.27809835],
  #                  dtype=np.float32).reshape(1, 1, 3)
  mean = np.array([0, 0, 0],
                   dtype=np.float32).reshape(1, 1, 3)
  std  = np.array([1, 1, 1],
                   dtype=np.float32).reshape(1, 1, 3)


  def __init__(self, opt, split):
    super(DOTA, self).__init__()
    self.data_dir = os.path.join(opt.data_dir, 'dota')
    self.img_dir = os.path.join(self.data_dir, '{}'.format(split))
    #self.annot_path = os.path.join(self.data_dir, "DOTA_train.json")

    self.annot_path = os.path.join(self.data_dir, 'DOTA_{}.json').format(split)

    self.max_objs = opt.K
    self.class_name = ['__background__', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J']
    #self.class_name = ['__background__', 'plane', 'ship', 'storage-tank', 'baseball-diamond', 'tennis-court', 'basketball-court', 'ground-track-field', 'harbor', 'bridge', 'large-vehicle', 'small-vehicle', 'helicopter', 'roundabout', 'soccer-ball-field', 'swimming-pool']
    self._valid_ids = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
    self.cat_ids = {v: i for i, v in enumerate(self._valid_ids)}
    self.voc_color = [(v // 32 * 64 + 64, (v // 8) % 4 * 64, v % 8 * 32) \
                      for v in range(1, self.num_classes + 1)]
    self._data_rng = np.random.RandomState(123)
    self._eig_val = np.array([0.2141788, 0.01817699, 0.00341571],
                             dtype=np.float32)
    self._eig_vec = np.array([
        [-0.58752847, -0.69563484, 0.41340352],
        [-0.5832747, 0.00994535, -0.81221408],
        [-0.56089297, 0.71832671, 0.41158938]
    ], dtype=np.float32)
    # self.mean = np.array([0.485, 0.456, 0.406], np.float32).reshape(1, 1, 3)
    # self.std = np.array([0.229, 0.224, 0.225], np.float32).reshape(1, 1, 3)

    self.split = split
    self.opt = opt

    print('==> initializing dota {} data.'.format(split))
    self.coco = coco.COCO(self.annot_path)
    self.images = self.coco.getImgIds()
    self.num_samples = len(self.images)

    print('Loaded {} {} samples'.format(split, self.num_samples))

  def _to_float(self, x):
    return float("{:.2f}".format(x))

  def convert_eval_format(self, all_bboxes):
    # import pdb; pdb.set_trace()
    detections = []
    for image_id in all_bboxes:
      for cls_ind in all_bboxes[image_id]:
        category_id = self._valid_ids[cls_ind - 1]
        for bbox in all_bboxes[image_id][cls_ind]:
          bbox[2] -= bbox[0]
          bbox[3] -= bbox[1]
          score = bbox[4]
          bbox_out  = list(map(self._to_float, bbox[0:4]))

          detection = {
              "image_id": int(image_id),
              "category_id": int(category_id),
              "bbox": bbox_out,
              "score": float("{:.2f}".format(score))
          }
          if len(bbox) > 5:
              extreme_points = list(map(self._to_float, bbox[5:13]))
              detection["extreme_points"] = extreme_points
          detections.append(detection)
    return detections

  def __len__(self):
    return self.num_samples

  def save_results(self, results, save_dir):
    json.dump(self.convert_eval_format(results), 
                open('{}/results.json'.format(save_dir), 'w'))
  

  def save_results_dota(self, results, save_dir):
    detections = []
    for image_id in results:
      for cls_ind in results[image_id]:
        category_id = self._valid_ids[cls_ind - 1]
        for bbox in results[image_id][cls_ind]:
          #bbox[2] -= bbox[0]
          #bbox[3] -= bbox[1]
          score = bbox[4]
          bbox_out = list(map(self._to_float, bbox[0:4]))
          image_name = self.coco.loadImgs(ids=[image_id])[0]['file_name']

          detection = {
            "image_id": int(image_id),
            "image_name": image_name,
            "category_id": int(category_id),
            "category_name": self.class_name[category_id],
            "bbox": bbox_out,
            "score": float("{:.4f}".format(score))
          }
          #if len(bbox) > 5:
          #  extreme_points = list(map(self._to_float, bbox[5:13]))
          #  detection["extreme_points"] = extreme_points
          detections.append(detection)
    if not os.path.exists(os.path.join(save_dir, 'result_dota')):
      os.makedirs(os.path.join(save_dir, 'result_dota'))
    #for cat_id in range(1, self.num_classes+1):
    for cat_id in range(1, self.num_classes+1):
      with open('%s/result_dota/Task2_%s.txt' % (save_dir,self.class_name[cat_id]), 'w') as file_writer:
        for detect_temp in detections:
          if detect_temp['category_id'] == cat_id:
            file_writer.write("%s %s %s %s %s %s\n" % (detect_temp['image_name'].split('.')[0], detect_temp['score'],
                                detect_temp['bbox'][0], detect_temp['bbox'][1], detect_temp['bbox'][2], detect_temp['bbox'][3]))
    return detections


  def save_results_plane(self, results, save_dir, result_conf_thresh = 0.1):
    detections = []
    for image_id in results:
      for cls_ind in results[image_id]:
        category_id = self._valid_ids[cls_ind - 1]
        for bbox in results[image_id][cls_ind]:
          #bbox[2] -= bbox[0]
          #bbox[3] -= bbox[1]
          score = bbox[4]
          bbox_out = list(map(self._to_float, bbox[0:4]))
          image_name = self.coco.loadImgs(ids=[image_id])[0]['file_name']

          detection = {
            "image_id": int(image_id),
            "image_name": image_name,
            "category_id": int(category_id),
            "category_name": self.class_name[category_id],
            "bbox": bbox_out,
            "score": float("{:.4f}".format(score))
          }
          #if len(bbox) > 5:
          #  extreme_points = list(map(self._to_float, bbox[5:13]))
          #  detection["extreme_points"] = extreme_points
          detections.append(detection)
    if not os.path.exists(os.path.join(save_dir, 'result_plane')):
      os.makedirs(os.path.join(save_dir, 'result_plane'))
    image_name_lists = set([d["image_name"] for d in detections])
    for name_tmp in image_name_lists:
      detect_temp = [d for d in detections if d["image_name"] == name_tmp and d["score"] > result_conf_thresh]
      with open('%s/result_plane/%s.txt' % (save_dir, name_tmp.split('.')[0]), 'w') as file_writer:
        for d_temp in detect_temp:
          file_writer.write("%s %s %s %s %s %s\n" % (d_temp["category_name"], d_temp["score"],
                                                     d_temp["bbox"][0], d_temp["bbox"][1],
                                                     d_temp["bbox"][2], d_temp["bbox"][3]))

    return detections

  def run_eval(self, results, save_dir):
    # result_json = os.path.join(save_dir, "results.json")
    # detections  = self.convert_eval_format(results)
    # json.dump(detections, open(result_json, "w"))
    
    self.save_results_dota(results, save_dir)
    #self.save_results(results, save_dir)
    
    #coco_dets = self.coco.loadRes('{}/results.json'.format(save_dir))
    #coco_eval = COCOeval(self.coco, coco_dets, "bbox")
    #coco_eval.evaluate()
    #coco_eval.accumulate()
    #coco_eval.summarize()
    
